/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/*!
 * \file resize_nearest_bilinear_2d.cc
 * \brief
 */
#include "runtime/resize_nearest_bilinear_2d.h"
#include "op_util.h"

namespace optiling {

constexpr size_t IDX_0 = 0;
constexpr size_t IDX_1 = 1;
constexpr size_t IDX_2 = 2;
constexpr size_t IDX_3 = 3;

static std::string PrintTilingData(const ResizeNB2DInfoR* tilingDataPtr) {
  std::string tilingStr;

  tilingStr += std::to_string(tilingDataPtr->tilingKey) + ",";
  tilingStr += std::to_string(tilingDataPtr->subKey) + ",";
  tilingStr += std::to_string(tilingDataPtr->ubOffset) + ",";
  tilingStr += std::to_string(tilingDataPtr->usedCoreCnt) + ",";
  tilingStr += std::to_string(tilingDataPtr->mcPos) + ",";
  tilingStr += std::to_string(tilingDataPtr->coreInOffset) + ",";
  tilingStr += std::to_string(tilingDataPtr->coreOutOffset) + ",";
  tilingStr += std::to_string(tilingDataPtr->ncLoopUnit) + ",";
  tilingStr += std::to_string(tilingDataPtr->hLoopUnit) + ",";
  tilingStr += std::to_string(tilingDataPtr->wLoopUnit) + ",";
  tilingStr += std::to_string(tilingDataPtr->ncPerCore) + ",";
  tilingStr += std::to_string(tilingDataPtr->hPerCore) + ",";
  tilingStr += std::to_string(tilingDataPtr->wPerCore) + ",";
  tilingStr += std::to_string(tilingDataPtr->ncSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->inHSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->inWSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->outHSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->outWSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->inSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->outSize) + ",";
  tilingStr += std::to_string(tilingDataPtr->ubOffset2) + ",";
  tilingStr += std::to_string(tilingDataPtr->inWAlign) + ",";
  tilingStr += std::to_string(tilingDataPtr->outWAlign);

  return tilingStr;
}

static bool GetInOutShape(const gert::TilingContext* context, ResizeNB2DInfoR* tilingDataPtr) {
  auto rtInShape = context->GetInputShape(0);
  OPS_CHECK_NULL_WITH_CONTEXT(context, rtInShape);
  auto rtOutShape = context->GetOutputShape(0);
  OPS_CHECK_NULL_WITH_CONTEXT(context, rtOutShape);
  const gert::Shape& imageShape = rtInShape->GetStorageShape();
  const gert::Shape& yShape = rtOutShape->GetStorageShape();

  // the format is NCHW
  tilingDataPtr->ncSize = imageShape[IDX_0] * imageShape[IDX_1];
  tilingDataPtr->inHSize = imageShape[IDX_2];
  tilingDataPtr->inWSize = imageShape[IDX_3];
  tilingDataPtr->outHSize = yShape[IDX_2];
  tilingDataPtr->outWSize = yShape[IDX_3];

  return true;
}

static void CalcTilingKey(const gert::TilingContext* context, ResizeNB2DInfoR* tilingDataPtr,
                          const ResizeNB2DCompileInfo* compileInfo) {
  OP_LOGD(context->GetNodeName(), "begin to count tiling mode");
  if ((tilingDataPtr->inWSize == tilingDataPtr->outWSize) && (tilingDataPtr->inHSize == tilingDataPtr->outHSize)) {
    tilingDataPtr->tilingKey = TILING_NO_SCALING;
    return;
  }

  if ((tilingDataPtr->inWSize == 1) && (tilingDataPtr->inHSize == 1)) {
    tilingDataPtr->tilingKey = TILING_ONE_2_MANY;
    return;
  }

  if ((tilingDataPtr->outWSize == 1) && (tilingDataPtr->outHSize == 1)) {
    tilingDataPtr->tilingKey = TILING_MANY_2_ONE;
    return;
  }

  if ((tilingDataPtr->inHSize <= 1000) && (tilingDataPtr->outHSize <= 1000) &&
      (tilingDataPtr->inWSize <= 64) && (tilingDataPtr->outWSize <= 32)) {
    tilingDataPtr->tilingKey = TILING_SMALL_W;
    return;
  }

  // default
  tilingDataPtr->tilingKey = TILING_DEFAULT;
}

static bool SetMCTiling4NoScaling(const ResizeNB2DCompileInfo* compileInfo, const int64_t elePerBlock,
                                  ResizeNB2DInfoR* tilingDataPtr) {
  int64_t minLoopUnit = MIN_BLOCK_NUM * elePerBlock;
  int64_t nchwSize = tilingDataPtr->ncSize * tilingDataPtr->inWSize * tilingDataPtr->inHSize;
  int64_t coreNum = compileInfo->coreNum;

  // to make sure the elements in per core is block align or all
  tilingDataPtr->coreInOffset = std::min(ops::CeilAlign(ops::CeilDiv(nchwSize, coreNum), minLoopUnit), nchwSize);
  tilingDataPtr->usedCoreCnt = ops::CeilDiv(nchwSize, tilingDataPtr->coreInOffset);
  tilingDataPtr->coreOutOffset = tilingDataPtr->coreInOffset;
  tilingDataPtr->wLoopUnit = compileInfo->ubMaxNum;
  tilingDataPtr->inWSize = nchwSize;

  return true;
}

static bool SetMCTiling4One2Many(const ResizeNB2DCompileInfo* compileInfo, const int64_t elePerBlock,
                                 ResizeNB2DInfoR* tilingDataPtr) {
  int64_t minWLoopUnit = 32;
  int64_t minNCLoopUnit = 2 * elePerBlock;

  tilingDataPtr->outWSize = tilingDataPtr->outHSize * tilingDataPtr->outWSize;
  tilingDataPtr->outHSize = 1;
  tilingDataPtr->ubOffset = compileInfo->ubMaxNum / 2 / elePerBlock * elePerBlock;

  if (tilingDataPtr->ncSize < tilingDataPtr->outWSize / minWLoopUnit) {
    tilingDataPtr->mcPos = POS_HW;
    tilingDataPtr->subKey = 0;
    tilingDataPtr->ncPerCore = tilingDataPtr->ncSize;
    tilingDataPtr->ncLoopUnit = tilingDataPtr->ubOffset;
    tilingDataPtr->coreInOffset = 0;
    tilingDataPtr->coreOutOffset =
        std::min(ops::CeilAlign(ops::CeilDiv(tilingDataPtr->outWSize, compileInfo->coreNum), minWLoopUnit),
                 tilingDataPtr->outWSize);
    tilingDataPtr->usedCoreCnt = ops::CeilDiv(tilingDataPtr->outWSize, tilingDataPtr->coreOutOffset);
    tilingDataPtr->wPerCore = tilingDataPtr->coreOutOffset;
    tilingDataPtr->wLoopUnit = VDP_LEN;
  } else {
    // to keep nc per core is block align or all
    if (tilingDataPtr->outWSize <= minWLoopUnit) {
      tilingDataPtr->subKey = 1;
      tilingDataPtr->coreInOffset =
          std::min(ops::CeilAlign(ops::CeilDiv(tilingDataPtr->ncSize, compileInfo->coreNum), minNCLoopUnit),
                   tilingDataPtr->ncSize);
      tilingDataPtr->wLoopUnit = minWLoopUnit;
      tilingDataPtr->ncLoopUnit = tilingDataPtr->ubOffset / tilingDataPtr->outWSize / elePerBlock * elePerBlock;
    } else {
      tilingDataPtr->subKey = 0;
      tilingDataPtr->coreInOffset =
          std::min(ops::CeilDiv(tilingDataPtr->ncSize, compileInfo->coreNum), tilingDataPtr->ncSize);
      tilingDataPtr->wLoopUnit = VDP_LEN;
      tilingDataPtr->ncLoopUnit = tilingDataPtr->ubOffset;
    }
    tilingDataPtr->mcPos = POS_NC;
    tilingDataPtr->wPerCore = tilingDataPtr->outWSize;
    tilingDataPtr->coreOutOffset = tilingDataPtr->coreInOffset * tilingDataPtr->outWSize;
    tilingDataPtr->ncPerCore = tilingDataPtr->coreInOffset;
    tilingDataPtr->usedCoreCnt = ops::CeilDiv(tilingDataPtr->ncSize, tilingDataPtr->coreInOffset);
  }

  return true;
}

static bool SetMCTiling4Many2One(const ResizeNB2DCompileInfo* compileInfo, const int64_t elePerBlock,
                                 ResizeNB2DInfoR* tilingDataPtr) {
  int64_t hwSize = tilingDataPtr->inHSize * tilingDataPtr->inWSize;
  int64_t vgatherGate = MIN_BLOCK_NUM * elePerBlock;

  tilingDataPtr->ubOffset = compileInfo->ubMaxNum / 2 / elePerBlock * elePerBlock;
  if (hwSize <= vgatherGate) {
    tilingDataPtr->subKey = 0;
    tilingDataPtr->ncLoopUnit = tilingDataPtr->ubOffset / hwSize / elePerBlock * elePerBlock;
  } else {
    tilingDataPtr->subKey = 1;
    tilingDataPtr->ncLoopUnit = tilingDataPtr->ubOffset / elePerBlock / elePerBlock * elePerBlock;
  }
  tilingDataPtr->ncPerCore = std::min(
      ops::CeilAlign(ops::CeilDiv(tilingDataPtr->ncSize, compileInfo->coreNum), elePerBlock), tilingDataPtr->ncSize);
  tilingDataPtr->usedCoreCnt = ops::CeilDiv(tilingDataPtr->ncSize, tilingDataPtr->ncPerCore);
  tilingDataPtr->coreOutOffset = tilingDataPtr->ncPerCore;
  tilingDataPtr->coreInOffset = tilingDataPtr->coreOutOffset * hwSize;

  return true;
}

static void SetTilingCoreParam4SplitNC(ResizeNB2DInfoR* tilingDataPtr, int64_t coreNum) {
  tilingDataPtr->ncPerCore = ops::CeilDiv(tilingDataPtr->ncSize, coreNum);
  tilingDataPtr->usedCoreCnt = ops::CeilDiv(tilingDataPtr->ncSize, tilingDataPtr->ncPerCore);
  tilingDataPtr->inSize = tilingDataPtr->inHSize * tilingDataPtr->inWSize;
  tilingDataPtr->outSize = tilingDataPtr->outHSize * tilingDataPtr->outWSize;
  tilingDataPtr->coreInOffset = tilingDataPtr->ncPerCore * tilingDataPtr->inSize;
  tilingDataPtr->coreOutOffset = tilingDataPtr->ncPerCore * tilingDataPtr->outSize;
}

static bool SetMCTiling4SmallW(const ResizeNB2DCompileInfo* compileInfo, const int64_t elePerBlock,
                               ResizeNB2DInfoR* tilingDataPtr) {
  SetTilingCoreParam4SplitNC(tilingDataPtr, compileInfo->coreNum);

  tilingDataPtr->hLoopUnit = UB_IDX_SIZE - tilingDataPtr->outWAlign * INDEX_DATA_SIZE;

  tilingDataPtr->outWAlign = ops::CeilAlign(tilingDataPtr->outWSize, elePerBlock);
  const int64_t inSizeAlign = ops::CeilAlign(tilingDataPtr->inSize, elePerBlock);
  tilingDataPtr->ncLoopUnit = std::min(ops::FloorDiv(compileInfo->ubMaxNum,
                                                     inSizeAlign + tilingDataPtr->outWAlign * tilingDataPtr->outHSize),
                                       tilingDataPtr->ncPerCore);
  if (tilingDataPtr->ncLoopUnit == 0) {
     // 先不考虑需要分H轴的情况，这里没写完
    tilingDataPtr->mcPos = POS_HW;
    tilingDataPtr->ncLoopUnit = 1;
  } else {
    tilingDataPtr->mcPos = POS_NC;
    tilingDataPtr->ubOffset = inSizeAlign * tilingDataPtr->ncLoopUnit;
  }

  return true;
}

static bool SetMCTiling4Default(const ResizeNB2DCompileInfo* compileInfo, const int64_t elePerBlock,
                                ResizeNB2DInfoR* tilingDataPtr) {
  SetTilingCoreParam4SplitNC(tilingDataPtr, compileInfo->coreNum);

  constexpr int64_t W_SUBKEY_FACTOR = 10;
  constexpr int64_t H_SUBKEY_FACTOR = 100;
  constexpr int64_t SUBKEY_EQUAL = 1;
  constexpr int64_t SUBKEY_ONE_2_MANY = 2;
  constexpr int64_t SUBKEY_MANY_2_ONE = 3;
  constexpr int64_t SUBKEY_INT_SCALE_ZOOM_OUT = 4;
  constexpr int64_t SUBKEY_INT_SCALE_ZOOM_IN = 5;
  if (tilingDataPtr->inWSize == tilingDataPtr->outWSize) {
    tilingDataPtr->subKey += W_SUBKEY_FACTOR * SUBKEY_EQUAL;
  } else if (tilingDataPtr->inWSize == 1) {
    tilingDataPtr->subKey += W_SUBKEY_FACTOR * SUBKEY_ONE_2_MANY;
  } else if (tilingDataPtr->outWSize == 1) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_MANY_2_ONE;
  } else if (tilingDataPtr->outWSize > tilingDataPtr->inWSize && tilingDataPtr->outWSize % tilingDataPtr->inWSize == 0) {
    tilingDataPtr->subKey += W_SUBKEY_FACTOR * SUBKEY_INT_SCALE_ZOOM_OUT;
  } else if (tilingDataPtr->inWSize > tilingDataPtr->outWSize && tilingDataPtr->inWSize % tilingDataPtr->outWSize == 0) {
    tilingDataPtr->subKey += W_SUBKEY_FACTOR * SUBKEY_INT_SCALE_ZOOM_IN;
  }
  if (tilingDataPtr->inHSize == tilingDataPtr->outHSize) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_EQUAL;
  } else if (tilingDataPtr->inHSize == 1) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_ONE_2_MANY;
  } else if (tilingDataPtr->outHSize == 1) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_MANY_2_ONE;
  } else if (tilingDataPtr->outHSize > tilingDataPtr->inHSize && tilingDataPtr->outHSize % tilingDataPtr->inHSize == 0) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_INT_SCALE_ZOOM_OUT;
  } else if (tilingDataPtr->inHSize > tilingDataPtr->outHSize && tilingDataPtr->inHSize % tilingDataPtr->outHSize == 0) {
    tilingDataPtr->subKey += H_SUBKEY_FACTOR * SUBKEY_INT_SCALE_ZOOM_IN;
  }

  tilingDataPtr->hLoopUnit = UB_IDX_SIZE;

  const int64_t eleSize = BYTES_PER_BLOCK / elePerBlock;
  const int64_t ubSize = compileInfo->ubMaxNum * eleSize;
  tilingDataPtr->inWAlign = ops::CeilAlign(tilingDataPtr->inWSize, elePerBlock);
  tilingDataPtr->outWAlign = ops::CeilAlign(tilingDataPtr->outWSize, elePerBlock);
  const int64_t inWUbSize = tilingDataPtr->inWAlign * eleSize;
  const int64_t outWUbSize = tilingDataPtr->outWAlign * eleSize;
  const int64_t outWUbIdxSize = tilingDataPtr->outWAlign * INDEX_DATA_SIZE;
  if (ubSize >= inWUbSize + outWUbSize + outWUbIdxSize) {
    tilingDataPtr->mcPos = POS_NC;
    tilingDataPtr->wLoopUnit = tilingDataPtr->outWSize;
    tilingDataPtr->ncLoopUnit = ops::FloorDiv(ubSize - outWUbIdxSize,
                                              (tilingDataPtr->inWAlign + tilingDataPtr->outWAlign) * eleSize);
  } else { // 先不不考虑W过大的情况，有可能用别的模板
    tilingDataPtr->mcPos = POS_HW;
    tilingDataPtr->ncLoopUnit = 1;
  }
  tilingDataPtr->ubOffset = tilingDataPtr->ncLoopUnit * tilingDataPtr->inWAlign;
  tilingDataPtr->ubOffset2 = tilingDataPtr->ubOffset + tilingDataPtr->ncLoopUnit * tilingDataPtr->outWAlign;

  return true;
}

static bool Set2DMultiCoreTilingParas(const gert::TilingContext* context, const ResizeNB2DCompileInfo* compileInfo,
                                      ResizeNB2DInfoR* tilingDataPtr) {
  auto dType = context->GetInputDesc(IDX_0)->GetDataType();
  int64_t dTypeSize = ge::GetSizeByDataType(dType);
  int64_t elePerBlock = 16;

  if (dTypeSize > 0) {
    elePerBlock = BYTES_PER_BLOCK / dTypeSize;
  }

  if (tilingDataPtr->tilingKey == TILING_NO_SCALING) {
    OP_TILING_CHECK(!SetMCTiling4NoScaling(compileInfo, elePerBlock, tilingDataPtr),
                    VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "set parameters for no scaling failed"),
                    return false);
  } else if (tilingDataPtr->tilingKey == TILING_ONE_2_MANY) {
    OP_TILING_CHECK(!SetMCTiling4One2Many(compileInfo, elePerBlock, tilingDataPtr),
                    VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "set parameters for one to many failed"),
                    return false);
  } else if (tilingDataPtr->tilingKey == TILING_MANY_2_ONE) {
    OP_TILING_CHECK(!SetMCTiling4Many2One(compileInfo, elePerBlock, tilingDataPtr),
                    VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "set parameters for many to one failed"),
                    return false);
  } else if (tilingDataPtr->tilingKey == TILING_SMALL_W) {
    OP_TILING_CHECK(!SetMCTiling4SmallW(compileInfo, elePerBlock, tilingDataPtr),
                    VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "set parameters for many to one failed"),
                    return false);
  } else {
    OP_TILING_CHECK(!SetMCTiling4Default(compileInfo, elePerBlock, tilingDataPtr),
                    VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "set parameters for many to one failed"),
                    return false);
  }

  return true;
}

static void InitTilingData(ResizeNB2DInfoR* tilingDataPtr) {
  tilingDataPtr->tilingKey = 0;
  tilingDataPtr->subKey = 0;
  tilingDataPtr->ubOffset = 0;
  tilingDataPtr->usedCoreCnt = 1;
  tilingDataPtr->mcPos = 0;
  tilingDataPtr->coreInOffset = 0;
  tilingDataPtr->coreOutOffset = 0;
  tilingDataPtr->ncLoopUnit = 1;
  tilingDataPtr->hLoopUnit = 1;
  tilingDataPtr->wLoopUnit = 1;
  tilingDataPtr->ncPerCore = 1;
  tilingDataPtr->hPerCore = 1;
  tilingDataPtr->wPerCore = 1;
  tilingDataPtr->ncSize = 1;
  tilingDataPtr->inHSize = 1;
  tilingDataPtr->inWSize = 1;
  tilingDataPtr->outHSize = 1;
  tilingDataPtr->outWSize = 1;
  tilingDataPtr->inSize = 1;
  tilingDataPtr->outSize = 1;
  tilingDataPtr->ubOffset2 = 0;
  tilingDataPtr->inWAlign = 1;
  tilingDataPtr->outWAlign = 1;
}

ge::graphStatus TilingForResizeNB2D(gert::TilingContext* context) {
  OP_LOGD(context->GetNodeName(), "ResizeNB2D tiling running begin");
  const ResizeNB2DCompileInfo* compileInfo = reinterpret_cast<const ResizeNB2DCompileInfo*>(context->GetCompileInfo());
  OPS_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
  ResizeNB2DInfoR* tilingDataPtr = context->GetTilingData<ResizeNB2DInfoR>();
  OPS_CHECK_NULL_WITH_CONTEXT(context, tilingDataPtr);
  InitTilingData(tilingDataPtr);

  OP_TILING_CHECK(!GetInOutShape(context, tilingDataPtr),
                  VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get input or output shape failed"),
                  return ge::GRAPH_FAILED);
  CalcTilingKey(context, tilingDataPtr, compileInfo);
  OP_TILING_CHECK(
      !Set2DMultiCoreTilingParas(context, compileInfo, tilingDataPtr),
      VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "count multiple core tiling parameters failed"),
      return ge::GRAPH_FAILED);

  context->SetBlockDim(static_cast<uint32_t>(tilingDataPtr->usedCoreCnt));

  OP_LOGI(context->GetNodeName(), "the tiling data is: %s.", PrintTilingData(tilingDataPtr).c_str());

  return ge::GRAPH_SUCCESS;
}

ge::graphStatus TilingPrepareForResizeNB2D(gert::TilingParseContext* context) {
  auto compileInfo = GetCompileInfoPtr<ResizeNB2DCompileInfo>(context);
  OPS_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
  std::unique_ptr<nlohmann::json> parsedObjectCInfo = GetCompileInfoJson(context);
  OPS_CHECK_NULL_WITH_CONTEXT(context, parsedObjectCInfo);

  const nlohmann::json& vars = (*parsedObjectCInfo)["vars"];
  OP_TILING_CHECK(vars.empty(), VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get vars failed."),
                  return ge::GRAPH_FAILED);

  uint32_t coreNumVar = 0;
  OP_TILING_CHECK(!GetTilingCoreNum(context, coreNumVar),
                  VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get core info from platform faided."),
                  return ge::GRAPH_FAILED);
  compileInfo->coreNum = static_cast<int64_t>(coreNumVar);
  OP_TILING_CHECK(!ReadCompileItem(vars, "ub_max_num", compileInfo->ubMaxNum),
                  VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get ub_ele from compile info faided."),
                  return ge::GRAPH_FAILED);
  OP_TILING_CHECK(
      !ReadCompileItem(vars, "align_corners", compileInfo->alignCorners),
      VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get align_corners from compile info faided."),
      return ge::GRAPH_FAILED);
  OP_TILING_CHECK(
      !ReadCompileItem(vars, "half_pixel_centers", compileInfo->halfPixelCenters),
      VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get half_pixel_centers from compile info faided."),
      return ge::GRAPH_FAILED);
  OP_TILING_CHECK(!ReadCompileItem(vars, "resize_mode", compileInfo->resizeMode),
                  VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "get resize mode from compile info faided."),
                  return ge::GRAPH_FAILED);

  return ge::GRAPH_SUCCESS;
}
}  // namespace optiling
